from fastapi import FastAPI
import uvicorn
import requests
import pydicom  
from io import BytesIO
import numpy as np
import onnxruntime
import time
from pydantic import BaseModel
import logging
from logging.handlers import RotatingFileHandler
from collections import Counter
from typing import Tuple, Optional

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("fastapi_logger")

# 设置日志文件和日志轮转
file_handler = RotatingFileHandler("app.log", maxBytes=10485760, backupCount=5)
file_handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
logger.addHandler(file_handler)

app = FastAPI()

class Item(BaseModel):
    url: str

label_dict = {
        "左冠足位": 'F',
        "左冠左肩位": 'LS',
        "右冠头位": 'RCA',
        "右冠左前斜位(偏头位)": 'RCA',
        "左冠头位": 'other',
        "左冠蜘蛛位": 'other',
        "左冠右肩位": 'other',
        "左冠肝位": 'other',
        "": 'other',
    }

def softmax(x):
    # 计算每行的最大值
    row_max = x.max(axis=1).reshape(-1, 1)
    # 从每行中减去最大值以确保数值稳定
    x_exp = np.exp(x - row_max)
    # 每行求和
    x_sum = np.sum(x_exp, axis=1, keepdims=True)
    # 计算 softmax 值
    s = x_exp / x_sum
    return s


def classify_xa_image_angles(
    primary_angle: Optional[float], secondary_angle: Optional[float]) -> Tuple[Optional[str], Optional[str]]:
    """
    Classifies the XA image angles.

    :param primary_angle: Primary angle.
    :param secondary_angle: Secondary angle.
    :return: Tuple of classifications for primary and secondary angles.
    """
    if primary_angle is None or secondary_angle is None:
        return None, None

    if -15 <= primary_angle <= 15:
        primary_angle_class = "AP"
    elif 15 < primary_angle <= 90:
        primary_angle_class = "LAO"
    elif -90 <= primary_angle < -15:
        primary_angle_class = "RAO"
    else:
        primary_angle_class = "Unknown"

    secondary_angle_class = "CRA" if secondary_angle >= 0 else "CAU"

    return primary_angle_class, secondary_angle_class


def determine_machine_position(
    primary_angle: Optional[float],
    secondary_angle: Optional[float],
    primary_angle_class: Optional[str],
    secondary_angle_class: Optional[str],) -> Tuple[str, str]:
    """
    Determines the machine position based on the classified angles.

    :param primary_angle: Primary angle.
    :param secondary_angle: Secondary angle.
    :param primary_angle_class: Classification of the primary angle.
    :param secondary_angle_class: Classification of the secondary angle.
    :return: Tuple of positions for left and right coronary arteries.
    """
    if None in [
        primary_angle,
        secondary_angle,
        primary_angle_class,
        secondary_angle_class,
    ]:
        return "未知体位", "未知体位"

    positions = {
        ("AP", "CAU"): (
            ("左冠足位", "右冠头位") if -45 <= secondary_angle <= 0 else None
        ),
        ("AP", "CRA"): ("左冠头位", "右冠头位") if 0 < secondary_angle <= 45 else None,
        ("LAO", "CAU"): (
            ("左冠蜘蛛位", "右冠左前斜位(偏头位)")
            if 16 <= primary_angle <= 45 and -45 <= secondary_angle <= 0
            else None
        ),
        ("LAO", "CRA"): (
            ("左冠左肩位", "右冠左前斜位(偏头位)")
            if 16 <= primary_angle <= 45 and 0 < secondary_angle <= 45
            else None
        ),
        ("RAO", "CAU"): (
            ("左冠肝位", "")
            if -45 <= primary_angle <= -16 and -45 <= secondary_angle <= 0
            else None
        ),
        ("RAO", "CRA"): (
            ("左冠右肩位", "")
            if -45 <= primary_angle <= -16 and 0 < secondary_angle <= 45
            else None
        ),
    }

    return positions.get(
        (primary_angle_class, secondary_angle_class), ("未知体位", "未知体位")
    )

def pred_r_or_l_from_onnx(frames, data):
    providers=['CPUExecutionProvider']
    onnx_session = onnxruntime.InferenceSession('vessel_position_rca_v2.0.onnx', providers=providers)
    result_label = []
    result_pred = []
    for frame in frames:
        new_img = np.reshape(data[frame, :, :], (1, 1, 512,512))
        new_img = new_img.repeat(3, axis=1)
        # mean=np.array([123.675, 116.28, 103.53]).reshape((3,1,1))
        # std=np.array([58.395, 57.12, 57.375]).reshape((3,1,1))
        # new_img = (new_img - mean) / std
        new_img = (new_img - new_img.min()) / (new_img.max() - new_img.min())

        new_img = new_img.astype(np.float32)
        pred = onnx_session.run(None, {"input": new_img})
        predict = pred[0]
        max_label = np.argmax(predict)
        max_pred = softmax(predict)[0][max_label]
        result_label.append(max_label)
        result_pred.append(max_pred) 

    result_label_count = Counter(result_label)
    max_label_pred = result_label_count.most_common(1)[0][0]
    result_print_str = [result_pred[i] for i, category in enumerate(result_label) if category == max_label_pred]
    logger.info(result_print_str)
    average_pred_for_most_common = sum(result_print_str) / len(result_print_str)
    logger.info(average_pred_for_most_common)
    return max_label_pred, average_pred_for_most_common

@app.get('/')
def hello():
    return "hello world!!!"

@app.post('/inference/')
def inference(item: Item):
    start_time = time.time()

    try:
        # response = requests.get(url.message)  # 设置超时时间
        response = requests.get(item.url)  # 设置超时时间
        response.raise_for_status()  # 如果请求不是200 OK，会抛出HTTPError异常
        bytes_dcm = BytesIO(response.content)
    except requests.exceptions.RequestException as e:
        return {'error': 'dicom文件找不到'}
    dcm = pydicom.dcmread(bytes_dcm)
    logger.info('url : {}'.format(item.url))
    primary_angle = dcm.PositionerPrimaryAngle
    secondary_angle = dcm.PositionerSecondaryAngle    
    primary_angle_class, secondary_angle_class = classify_xa_image_angles(primary_angle, secondary_angle)
    result_angle = determine_machine_position(primary_angle, secondary_angle, primary_angle_class, secondary_angle_class)
    dcm_array = dcm.pixel_array
    if len(dcm_array.shape) != 3 or dcm_array.shape[-1] !=512:
        return {'error': 'dicom文件格式错误'}

    study_instance_uid = dcm.StudyInstanceUID
    series_instance_uid = dcm.SeriesInstanceUID
    sop_instance_uid = dcm.SOPInstanceUID
    data = dcm_array
    logger.info('time 1:{}'.format(time.time()-start_time))
    start_time = time.time()

    frames = []
    for num in np.linspace(0, data.shape[0]-1, 10)[1:-1]:
        frames.append(int(num)) 
    max_label_pred, average_pred_for_most_common = pred_r_or_l_from_onnx(frames, data)
    if result_angle == None:
        result_angle = ["",""]
    
    if max_label_pred == 0:
        position_result = label_dict[result_angle[0]]
    elif max_label_pred == 1:
        position_result = label_dict[result_angle[1]]
    else:
        position_result = 'unknown'

    return_item = {
        "study_instance_uid": study_instance_uid,
        "series_instance_uid": series_instance_uid,
        "sop_instance_uid": sop_instance_uid,
        "position": position_result,
        "value": str(round(average_pred_for_most_common, 2)),
    }

    logger.info('result :{}'.format(result_angle[max_label_pred]))
    
    logger.info('time :{}'.format(time.time()-start_time))
    return return_item

if __name__ == '__main__':
    uvicorn.run(app, host='192.168.1.71', port=8000,log_level="info")